###############################
#	prj: TBtools-plugin
#	Assignment: One step WGCNA
#	Author: Shawn Wang
#	Date: Jan 12, 2021
###############################
con <- file("test.log") 
sink(con, append=TRUE) 
sink(con, append=TRUE, type="message") 
##================Step1 packages and args===========================##
options(stringsAsFactors = F)
options("repos" = c(CRAN="https://mirrors.tuna.tsinghua.edu.cn/CRAN/"))
argv <- commandArgs(TRUE)
raw<- argv[1] #'expression data matrix'
traitData <- argv[2]# 'trait data'
datatype <- argv[3] # 'data type readcount or fpkm'
method <- argv[4] # 'Which standardization method do you prefer'
RcCutoff <- argv[5]# 'Threshold of readcount '
RcCutoff = as.numeric(RcCutoff)
samplePerc <- argv[6] # 'At least in how many samples readcount is greater than this threshold '
samplePerc = as.numeric(samplePerc)
RemainGeneNum <- argv[7] # 'How many genes do you want to use for WGCNA '
RemainGeneNum = as.numeric(RemainGeneNum)
title <- argv[8] #' output data'
########################
####   test      ######
# setwd("~/01.TBtoolsTest/")
# raw = "~/01.TBtoolsTest/BtJ.count.xls"
# traitData = "~/01.TBtoolsTest/trait.txt"
# datatype = "count"
# method = "lgcpm"
# RcCutoff = "0"
# samplePerc = "0"
# RemainGeneNum = "340"
# RcCutoff = as.numeric(RcCutoff)
# samplePerc = as.numeric(samplePerc)
# RemainGeneNum = as.numeric(RemainGeneNum)
# title = "test"
########################
Title = title
type = "unsigned"
corType = "pearson"
maxPOutliers = ifelse(corType=="pearson",1,0.05)
robustY = ifelse(corType=="pearson",T,F)

## packages
if (!requireNamespace("BiocManager", quietly = TRUE))
  install.packages("BiocManager")
if (!require('WGCNA')) BiocManager::install('GO.db',update = FALSE)
if (!require('WGCNA')) BiocManager::install('WGCNA',update = FALSE)
if (!require('ggplot2')) install.packages('ggplot2')
if (!require('dplyr')) install.packages('dplyr')
if (!require('stringr')) install.packages('stringr')
if (!require('DESeq2')) BiocManager::install('DESeq2',update = FALSE)
if (!require('ape')) install.packages('ape')
if (!require('edgeR')) BiocManager::install('edgeR',update = FALSE)
if (!require('reshape2')) install.packages('reshape2')
suppressMessages(library(DESeq2))
suppressMessages(library(ggplot2))
suppressMessages(library(dplyr))
suppressMessages(library(WGCNA))
suppressMessages(library(stringr))
suppressMessages(library(ape))
suppressMessages(library(reshape2))
suppressMessages(library(edgeR))
##================Step2 Exp data and trait data===========================##
rawdata = read.delim(raw, header = T,
                     sep = "\t")
rawdata <- data.frame(row.names = rawdata[,1],
                      rawdata[,-1])
## func1 count 2 DEseq2 normdata
countvarTran = function(rawcount,RcCutoff,samplePerc) {
  samnum <- ncol(rawcount)
  casenum = ceiling(samnum/2)
  controlnum = samnum - casenum
  condition <- factor(c(rep("case",casenum),rep("control",controlnum)),levels = c("case","control"))
  colData <- data.frame(row.names = colnames(rawcount), condition)
  ## remove background noise
  x <- rawcount[apply(rawcount,1,function(x) sum(x > RcCutoff) > (samplePerc*ncol(rawcount))),]
  
  ## readcount standardization by DESeq2
  dds <- DESeqDataSetFromMatrix(x, colData, design = ~ condition)
  dds <- DESeq(dds)
  vsd <- assay(varianceStabilizingTransformation(dds))
  datExpr = data.frame(vsd)
  return(datExpr)
}
## func2 count 2 cpm
countCPM = function(rawcount, RcCutoff,samplePerc) {
  x <- rawcount[apply(rawcount,1,function(x) sum(x > RcCutoff) > (samplePerc*ncol(rawcount))),]
  datExpr  <- log10(edgeR::cpm(x)+1)
  return(datExpr)
}
## raw fpkm filter
fpkmfilter = function(rawcount, RcCutoff,samplePerc) {
  x <- rawcount[apply(rawcount,1,function(x) sum(x > RcCutoff) > (samplePerc*ncol(rawcount))),]
  datExpr  <- x
  return(datExpr)
}

## log fpkm filter
lgfpkmfilter = function(rawcount, RcCutoff,samplePerc) {
  x <- rawcount[apply(rawcount,1,function(x) sum(x > RcCutoff) > (samplePerc*ncol(rawcount))),]
  datExpr  <- log10(x+1)
  return(datExpr)
}
## datExpr final
if (
  datatype == "count" & method == "varianceStabilizingTransformation"
) {
  datExpr = countvarTran(rawcount = rawdata,RcCutoff = RcCutoff, samplePerc = samplePerc)
} else if (
  datatype == "count" & method == "lgcpm"
) {
  datExpr = countCPM(rawcount = rawdata,RcCutoff = RcCutoff, samplePerc = samplePerc)
} else if (
  datatype == "FPKM" & method == "rawFPKM"
) {
  datExpr = fpkmfilter(rawcount = rawdata,RcCutoff = RcCutoff, samplePerc = samplePerc)
} else if (
  datatype == "FPKM" & method == "lgFPKM"
) {
  datExpr = lgfpkmfilter(rawcount = rawdata,RcCutoff = RcCutoff, samplePerc = samplePerc)
}

## trait data
trait <- read.delim(traitData,header = T,
                    sep = "\t")
if (ncol(trait) == 2) {
  x <- trait
  Tcol = as.character(unique(x[,2]))
  b <- list()
  for (i in 1:length(Tcol)) {
    b[[i]] = data.frame(row.names = x[,1],
                        levels = ifelse(x[,2] == Tcol[i],1,0))
  }
  c <- bind_cols(b)
  c <- data.frame(row.names = x$name,
                  c)
  colnames(c) = Tcol
  rownames(c) = trait[,1]
  pheTmp <- c
} else {
  pheTmp = data.frame(row.names = trait[,1],
                      trait[,-1])
}

##================Step3 SFT ===========================##
WGCNA.SFT <- function(datExpr, Title, GeneNumCut){
  #reference: Tong Chen- sheng xin bao dian, the orignal link of the jianshu is missing. sry...
  datExpr <- datExpr
  type = "unsigned"
  corType = "pearson"
  corFnc = cor
  maxPOutliers = ifelse(corType=="pearson",1,0.05)
  robustY = ifelse(corType=="pearson",T,F)
  m.mad <- apply(datExpr,1,mad)
  if(GeneNumCut == 0){
    datExprVar = datExpr
  } else {
    datExprVar <- datExpr[which(m.mad > 
                                  max(quantile(m.mad, probs=seq(0, 1, GeneNumCut))[2],0.01)),]
  }
  dim(datExprVar)
  datExpr <- as.data.frame(t(datExprVar))
  ## 检测缺失值
  gsg = goodSamplesGenes(datExpr, verbose = 3)
  if (!gsg$allOK){
    # Optionally, print the gene and sample names that were removed:
    if (sum(!gsg$goodGenes)>0) 
      printFlush(paste("Removing genes:", 
                       paste(names(datExpr)[!gsg$goodGenes], collapse = ",")));
    if (sum(!gsg$goodSamples)>0) 
      printFlush(paste("Removing samples:", 
                       paste(rownames(datExpr)[!gsg$goodSamples], collapse = ",")));
    # Remove the offending genes and samples from the data:
    datExpr = datExpr[gsg$goodSamples, gsg$goodGenes]
  }
  ## sample cluster based on expression values
  nGenes = ncol(datExpr)
  nSamples = nrow(datExpr)
  assign("nGenes",value = nGenes, envir = globalenv())
  assign("nSamples",value = nSamples, envir = globalenv())
  dim(datExpr)
  #head(datExpr)[,1:8]
  ## trait-sample tree
  sampleTree = hclust(dist(datExpr), method = "average")
  
  pdf(file = paste("01",Title,"Sample_clustering.pdf",sep = "."),width = 28,height = 5)
  plot(sampleTree, main = "Sample clustering to detect outliers", sub="", xlab="")
  dev.off();
  # sample_colors <- numbers2colors(as.numeric(factor(dataTraits$Group)),
  #                               colors = c("green","yellow","red"),
  #                               signed = FALSE)
  # pdf(file = paste(Title,"Trait_Sample_clustering.pdf",sep = "."),width = 28,height = 5)
  # par(mar = c(1,4,3,1),cex=0.8)
  # {plotDendroAndColors(sampleTree, sample_colors,
  #                   groupLabels = colnames(sample),
  #                   cex.dendroLabels = 0.8,
  #                   marAll = c(1, 4, 3, 1),
  #                   cex.rowText = 0.01,
  #                   main = "sample dedrogram and trait heatmap")
  # }
  # dev.off();
  # export tree as nwk
  mytree <- as.phylo(sampleTree)
  write.tree(mytree,file = paste("01",Title,"samplecluster.nwk",sep = "."))
  # sft plot
  powers = c(c(1:10), seq(from = 12, to=30, by=2))
  sft = pickSoftThreshold(datExpr, powerVector=powers, 
                          networkType=type, verbose=5)
  
  pdf(file = paste("01",Title,"SFTPlot.pdf",sep = "."),width = 10,height = 7)
  par(mfrow = c(1,2))
  cex1 = 0.9
  # 横轴是Soft threshold (power)，纵轴是无标度网络的评估参数，数值越高，
  # 网络越符合无标度特征 (non-scale)
  {plot(sft$fitIndices[,1], -sign(sft$fitIndices[,3])*sft$fitIndices[,2],
        xlab="Soft Threshold (power)",
        ylab="Scale Free Topology Model Fit,signed R^2",type="n",
        main = paste("Scale independence"))
    text(sft$fitIndices[,1], -sign(sft$fitIndices[,3])*sft$fitIndices[,2],
         labels=powers,cex=cex1,col="red")
    # 筛选标准。R-square=0.85
    abline(h=0.90,col="red")
    abline(h=0.85,col="green")
    # Soft threshold与平均连通性
    plot(sft$fitIndices[,1], sft$fitIndices[,5],
         xlab="Soft Threshold (power)",ylab="Mean Connectivity", type="n",
         main = paste("Mean connectivity"))
    text(sft$fitIndices[,1], sft$fitIndices[,5], labels=powers, 
         cex=cex1, col="red")
  }
  dev.off();
  power = sft$powerEstimate
  power
  assign("datExpr",value = datExpr, envir = globalenv())
  assign("power",value = power, envir = globalenv())
  assign("sampleTree", value = sampleTree, envir = globalenv())
}
GNC = 1-RemainGeneNum/nrow(datExpr) ## how many genes do you want to retain.
WGCNA.SFT(Title = title,
          datExpr = datExpr,
          GeneNumCut = GNC)
if (is.na(power)){
  power = ifelse(nSamples<20, ifelse(type == "unsigned", 9, 18),
                 ifelse(nSamples<30, ifelse(type == "unsigned", 8, 16),
                        ifelse(nSamples<40, ifelse(type == "unsigned", 7, 14),
                               ifelse(type == "unsigned", 6, 12))
                 )
  )
}

##================Step4 module ===========================##
WGCNA.oneStepNetWork <- function(Title){
  #######============一步法网络构建===========
  # power: 上一步计算的软阈值
  # maxBlockSize: 计算机能处理的最大模块的基因数量 (默认5000)；
  #  4G内存电脑可处理8000-10000个，16G内存电脑可以处理2万个，32G内存电脑可
  #  以处理3万个
  #  计算资源允许的情况下最好放在一个block里面。
  # corType: pearson or bicor
  # numericLabels: 返回数字而不是颜色作为模块的名字，后面可以再转换为颜色
  # saveTOMs：最耗费时间的计算，存储起来，供后续使用
  # mergeCutHeight: 合并模块的阈值，越大模块越少
  cor <- WGCNA::cor
  net = blockwiseModules(datExpr, power = power, maxBlockSize = nGenes,
                         TOMType = type, minModuleSize = 25,
                         reassignThreshold = 0, mergeCutHeight = 0.4,
                         numericLabels = TRUE, pamRespectsDendro = FALSE,
                         saveTOMs=TRUE, corType = corType, 
                         maxPOutliers=maxPOutliers, loadTOMs=TRUE,
                         #saveTOMFileBase = paste(Title,".tom",sep = ""),
                         verbose = 3)
  
  assign("net",value = net, envir = globalenv())
  ## 灰色的为**未分类**到模块的基因。
  # Convert labels to colors for plotting
  moduleLabels = net$colors
  assign("moduleLabels",value = moduleLabels, envir = globalenv())
  moduleColors = labels2colors(moduleLabels)
  x = table(moduleColors)
  print(x)
  assign("moduleColors",value = moduleColors, envir = globalenv())
  # Plot the dendrogram and the module colors underneath
  # 如果对结果不满意，还可以recutBlockwiseTrees，节省计算时间
  pdf(file = paste("02",Title,"Module.pdf",sep = "."),width = 6,height = 6)
  plotDendroAndColors(net$dendrograms[[1]], moduleColors[net$blockGenes[[1]]],
                      "Module colors",
                      dendroLabels = FALSE, hang = 0.03,
                      addGuide = TRUE, guideHang = 0.05)
  dev.off()
  # module eigengene, 可以绘制线图，作为每个模块的基因表达趋势的展示
  MEs = net$MEs
  ### 不需要重新计算，改下列名字就好
  ### 官方教程是重新计算的，起始可以不用这么麻烦
  MEs_col = MEs
  colnames(MEs_col) = paste0("ME", labels2colors(
    as.numeric(str_replace_all(colnames(MEs),"ME",""))))
  MEs_col = orderMEs(MEs_col)
  assign("MEs",value = MEs, envir = globalenv())
  assign("MEs_col",value = MEs_col, envir = globalenv())
  
  # 根据基因间表达量进行聚类所得到的各模块间的相关性图
  # marDendro/marHeatmap 设置下、左、上、右的边距
  # cor<-stats::cor
  if (ncol(MEs_col)>3) {
    pdf(file = paste("02",Title,"Eigeng_adja_heatmap.pdf",sep = "."),width = 7,height = 10)
    plotEigengeneNetworks(MEs_col, "Eigengene adjacency heatmap", 
                          marDendro = c(3,3,2,4),
                          marHeatmap = c(3,4,2,2), plotDendrograms = T, 
                          xLabelsAngle = 90)
    dev.off()
  }
 
  Gene2module <- data.frame(GID = colnames(datExpr),
                            Module = moduleColors)
  write.table(Gene2module,file = paste("02",Title,"Gene2module.xls",sep = "."),
              row.names = F,
              quote = F,
              sep = "\t")
}
WGCNA.oneStepNetWork(Title = Title)

##================Step4 module-trait ===========================##
WGCNA.ModuleTrait <- function(Title,phenotype){
  traitData <- phenotype
  dim(traitData)
  ### 模块与表型数据关联
  if (corType=="pearson") {
    modTraitCor = cor(MEs_col, traitData, use = "p")
    modTraitP = corPvalueStudent(modTraitCor, nSamples)
  } else {
    modTraitCorP = bicorAndPvalue(MEs_col, traitData, robustY=robustY)
    modTraitCor = modTraitCorP$bicor
    modTraitP   = modTraitCorP$p
  }
  
  ## Warning in bicor(x, y, use = use, ...): bicor: zero MAD in variable 'y'.
  ## Pearson correlation was used for individual columns with zero (or missing)
  ## MAD.
  
  # signif表示保留几位小数
  textMatrix = paste(signif(modTraitCor, 2), "\n(", signif(modTraitP, 1), ")", sep = "")
  dim(textMatrix) = dim(modTraitCor)
  pdf(file = paste("03",Title,"Module_trait.pdf",sep = "."),width = 20,height = 10)
  labeledHeatmap(Matrix = modTraitCor, xLabels = colnames(traitData), 
                 yLabels = colnames(MEs_col), 
                 cex.lab = 0.7, xLabelsAngle = 45, xLabelsAdj = 1,
                 ySymbols = substr(colnames(MEs_col),3,1000), colorLabels = FALSE, 
                 colors = blueWhiteRed(50), 
                 textMatrix = textMatrix, setStdMargins = FALSE, 
                 cex.text = 0.6, zlim = c(-1,1),
                 main = paste("Module-trait relationships"))
  dev.off()
}
phenotype = pheTmp
WGCNA.ModuleTrait(Title = Title,
                  phenotype = phenotype)
save(MEs_col,nSamples,corType,file = paste("00",Title,"datatraitbase.Rdata",sep = "."))
##================Step5 conectivity of each modular ===========================##
Title = title
connet=abs(cor(datExpr,use="p"))^6
Alldegrees1=intramodularConnectivity(connet, moduleColors)
datKME=signedKME(datExpr, MEs_col, outputColumnName="MM.")
write.table(datKME, paste("04",Title,"Conectivity_of_each_modular.xls",sep = "."),
            sep = "\t",
            row.names = T,
            quote = F)
save(... = datExpr, power,moduleColors,file = paste("00",Title,".Modular.Rdata",sep = ""))


sink()
sink(type="message")

cat(readLines("test.log"), sep="\n")
write.table(cat(readLines("test.log"), sep="\n"), "log.txt")
